import sys
import av
import time
import datetime
import traceback
import numpy as np
from utils import media_utils
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout,
                             QHBoxLayout, QPushButton, QLineEdit, QLabel, QListWidget,
                             QMessageBox, QGridLayout, QScrollArea)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer, QMutex
from PyQt5.QtGui import QImage, QPixmap


class StreamThread(QThread):
    frame_signal = pyqtSignal(str, QImage)
    stream_end_signal = pyqtSignal(str)  # 添加新的信号用于通知视频结束

    def __init__(self, cameras):
        super().__init__()
        self.cameras = cameras
        self.running = False
        self.mutex = QMutex()  # 添加互斥锁

    def run(self):
        self.running = True
        while self.running:
            self.mutex.lock()
            try:
                for camera in self.cameras:
                    if camera.inp_av:
                        try:
                            # 分别处理视频和音频
                            for packet in camera.inp_av.demux():
                                for frame in packet.decode():
                                    # 检查是否为视频帧
                                    if isinstance(frame, av.VideoFrame):
                                        image = frame.to_ndarray(format='rgb24')
                                        h, w = image.shape[:2]
                                        bytes_per_line = 3 * w
                                        q_image = QImage(image.data, w, h, bytes_per_line,
                                                         QImage.Format_RGB888).copy()
                                        self.frame_signal.emit(camera.name, q_image)

                                        if camera.is_recording:
                                            # 编码视频帧
                                            frame = av.VideoFrame.from_ndarray(image, format='rgb24')
                                            for packet in camera.video_stream.encode(frame):
                                                camera.out_av.mux(packet)

                                    # 检查是否为音频帧
                                    elif isinstance(frame,
                                                    av.AudioFrame) and camera.is_recording and camera.audio_stream:
                                        # 编码音频帧
                                        for packet in camera.audio_stream.encode(frame):
                                            camera.out_av.mux(packet)
                                break  # 每次循环只处理一个packet的帧，避免阻塞

                        except av.EOFError:
                            print(f"视频 {camera.name} 播放完成")
                            if camera.is_recording:
                                # 刷新剩余视频帧
                                for packet in camera.video_stream.encode(None):
                                    camera.out_av.mux(packet)
                                # 刷新剩余音频帧
                                if camera.audio_stream:
                                    for packet in camera.audio_stream.encode(None):
                                        camera.out_av.mux(packet)
                            camera.disconnect()
                            self.stream_end_signal.emit(camera.name)
                        except Exception as e:
                            traceback.print_exc()
                            print(f"流处理错误: {str(e)}")
            finally:
                self.mutex.unlock()
                time.sleep(0.001)

    def stop(self):
        self.running = False
        self.mutex.lock()
        try:
            for camera in self.cameras:
                if camera.is_recording:
                    try:
                        # 刷新剩余视频帧
                        for packet in camera.video_stream.encode(None):
                            camera.out_av.mux(packet)
                        # 刷新剩余音频帧
                        if camera.audio_stream:
                            for packet in camera.audio_stream.encode(None):
                                camera.out_av.mux(packet)
                        camera.stop_recording()
                    except Exception as e:
                        print(f"停止录制时出错: {str(e)}")
        finally:
            self.mutex.unlock()

        if not self.wait(3000):
            print("线程未能正常结束，强制终止")
            self.terminate()
            self.wait()


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.cameras = []
        self.stream_thread = None
        self.ui_mutex = QMutex()  # 添加UI互斥锁
        self.setup_ui()

    def setup_ui(self):
        self.setWindowTitle("多路RTSP视频录制工具")
        self.setGeometry(100, 100, 1200, 800)

        # 主布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        main_layout = QHBoxLayout(main_widget)

        # 左侧控制面板
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)
        left_panel.setFixedWidth(300)

        # 添加摄像头控件
        camera_name_label = QLabel("视频名称:")
        self.camera_name_input = QLineEdit("demo")
        camera_url_label = QLabel("视频地址:")
        # self.camera_url_input = QLineEdit("rtsp://admin:C2332416@192.168.2.35")
        self.camera_url_input = QLineEdit("test-video.mp4")

        add_button = QPushButton("添加摄像头")
        add_button.clicked.connect(self.add_camera)

        # 摄像头列表
        self.camera_list = QListWidget()
        delete_button = QPushButton("删除选中")
        delete_button.clicked.connect(self.delete_camera)

        # 控制按钮
        self.play_button = QPushButton("全部播放")
        self.play_button.clicked.connect(self.start_streaming)
        self.record_button = QPushButton("开始录制")
        self.record_button.clicked.connect(self.start_recording)
        self.stop_button = QPushButton("停止播放")
        self.stop_button.clicked.connect(self.stop_all)

        # 添加控件到左侧面板
        left_layout.addWidget(camera_name_label)
        left_layout.addWidget(self.camera_name_input)
        left_layout.addWidget(camera_url_label)
        left_layout.addWidget(self.camera_url_input)
        left_layout.addWidget(add_button)
        left_layout.addWidget(self.camera_list)
        left_layout.addWidget(delete_button)
        left_layout.addWidget(self.play_button)
        left_layout.addWidget(self.record_button)
        left_layout.addWidget(self.stop_button)
        left_layout.addStretch()

        # 右侧视频显示区域
        self.video_widget = QWidget()
        self.video_layout = QGridLayout(self.video_widget)

        scroll_area = QScrollArea()
        scroll_area.setWidget(self.video_widget)
        scroll_area.setWidgetResizable(True)

        # 添加左右面板到主布局
        main_layout.addWidget(left_panel)
        main_layout.addWidget(scroll_area)

    def add_camera(self):
        name = self.camera_name_input.text().strip()
        url = self.camera_url_input.text().strip()

        if not name or not url:
            QMessageBox.warning(self, "警告", "请输入摄像头名称和RTSP地址")
            return

        camera = media_utils.RTSPCamera(name, url)
        self.cameras.append(camera)
        self.camera_list.addItem(f"{name} - {url}")

        # 添加视频显示标签
        row = (len(self.cameras) - 1) // 2
        col = (len(self.cameras) - 1) % 2

        video_container = QWidget()
        video_layout = QVBoxLayout(video_container)
        name_label = QLabel(name)
        video_label = QLabel()
        video_label.setFixedSize(480, 360)
        video_layout.addWidget(name_label)
        video_layout.addWidget(video_label)

        self.video_layout.addWidget(video_container, row, col)

    def delete_camera(self):
        current_row = self.camera_list.currentRow()
        if current_row >= 0:
            camera = self.cameras[current_row]
            camera.disconnect()
            self.cameras.pop(current_row)
            self.camera_list.takeItem(current_row)

            # 重新布局视频显示区域
            self.clear_video_layout()
            self.rebuild_video_layout()

    def clear_video_layout(self):
        while self.video_layout.count():
            item = self.video_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()

    def rebuild_video_layout(self):
        for i, camera in enumerate(self.cameras):
            row = i // 2
            col = i % 2

            video_container = QWidget()
            video_layout = QVBoxLayout(video_container)
            name_label = QLabel(camera.name)
            video_label = QLabel()
            video_label.setFixedSize(480, 360)
            video_layout.addWidget(name_label)
            video_layout.addWidget(video_label)

            self.video_layout.addWidget(video_container, row, col)

    def start_streaming(self):
        self.ui_mutex.lock()
        try:
            for camera in self.cameras:
                if not camera.connect():
                    QMessageBox.warning(self, "错误", f"无法连接到摄像头: {camera.name}")
                    return

            if self.stream_thread is None:
                self.stream_thread = StreamThread(self.cameras)
                self.stream_thread.frame_signal.connect(self.update_frame)
                self.stream_thread.stream_end_signal.connect(self.handle_stream_end)  # 连接新的信号
                self.stream_thread.start()
        finally:
            self.ui_mutex.unlock()

    def start_recording(self):
        for camera in self.cameras:
            camera.start_recording()

    def stop_all(self):
        try:
            if self.stream_thread:
                self.stream_thread.stop()
                self.stream_thread = None

            for camera in self.cameras:
                try:
                    if camera.is_recording:
                        camera.stop_recording()
                    camera.disconnect()
                except Exception as e:
                    print(f"停止摄像头时出错: {str(e)}")
        except Exception as e:
            print(f"停止所有设备时出错: {str(e)}")

    def update_frame(self, camera_name, image):
        self.ui_mutex.lock()  # UI更新前加锁
        try:
            for i in range(self.video_layout.count()):
                container = self.video_layout.itemAt(i).widget()
                if container and container.layout().itemAt(0).widget().text() == camera_name:
                    video_label = container.layout().itemAt(1).widget()
                    # 使用 Qt.SmoothTransformation 提高缩放质量
                    scaled_image = image.scaled(video_label.size(),
                                                Qt.KeepAspectRatio,
                                                Qt.SmoothTransformation)
                    video_label.setPixmap(QPixmap.fromImage(scaled_image))
                    break
        finally:
            self.ui_mutex.unlock()  # 确保解锁

    def handle_stream_end(self, camera_name):
        """处理视频结束事件"""
        self.ui_mutex.lock()
        try:
            # 找到对应的摄像头
            for camera in self.cameras:
                if camera.name == camera_name:
                    camera.disconnect()
                    break

            # 检查是否所有视频都已结束
            all_ended = all(camera.inp_av is None for camera in self.cameras)
            if all_ended and self.stream_thread:
                self.stream_thread.stop()
                self.stream_thread = None
        finally:
            self.ui_mutex.unlock()

    def closeEvent(self, event):
        try:
            # 不使用UI互斥锁，避免死锁
            self.stop_all()
            event.accept()
        except Exception as e:
            print(f"关闭窗口时出错: {str(e)}")
            event.accept()  # 确保窗口能够关闭


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
